在发布几周后,OpenAI的ChatGPT仍在不断惊艳世界。GPT使用大规模预训练模型来降低AI在不同场景重新训练的成本。这一思路已经被广泛证明。在科学智能(AI for Science)领域中,虽然不同学科所关心的具体问题天差地别,但是万变不离其宗,最终都要回归于基本元素和量子物理原理,因此,在科学智能领域也具备诞生这样大规模预训练模型的条件。
近日,AI for Science领域最大的开源社区DeepModeling举办了2022年社区年会。会上,北京科学智能研究院(AI for Science Institute,Beijing)联合深势科技,发布了首个覆盖元素周期表近70种元素的深度势能原子间势函数预训练模型—— DPA-1。该成果由北京科学智能研究院、深势科技、北京应用物理与计算数学研究所共同研发,DPA-1 是"自然科学的GPT",可模拟原子规模高至100亿,目前已经在高性能合金、半导体材料设计等应用场景中证明了其领先性和优越性,并展现出了在人工智能领域难得可贵的“可解释性”潜力。这一突破也是AI for Science走向大规模工程化的重要里程碑。 大会现场,北京应用物理与计算数学研究所研究员王涵老师详细分享了DPA-1 预训练模型的相关工作。以下是演讲实录(在不改变原意的基础上,我们对内容进行了编辑整理)。
大家好!今天非常有幸有这样的一个机会,给大家汇报一下我们在 DP预训练模型方面的工作。这个工作是和张铎等同学一起完成的。
背景还是我们面临的非常经典的问题,就是势函数的建模问题。我们把它当做一个监督学习的这样的一个任务。我们的 label 是从密度泛函等第一性原理的这种电子结构的计算里面得到的。我们希望训练我们的势函数的模型,使得我们的势函数的模型在能量、受力的预测上面能够跟我们的 label 尽可能的接近。这就是我们的面临的势函数建模的这样一个监督学习的任务。Deep Potential 模型是在 2017 年提出的,到现在已经五年了,其实经过了很多的发展。我们也能够看到它帮助我们解决了很多在科学上非常有挑战的问题,比如像水的相图计算。但是今天其实并不想过多的谈论成绩,而是想谈论我们在使用 DeepPotential 时候面临的一些挑战。
当然挑战是非常多的,今天集中于其中的一个挑战——数据模型和应用方面的挑战。我们先来看看挑战是什么。首先,在数据这一端,我们有非常多的数据的来源,包括 DP Library(即现在的 AIS-Square)上面有很多的数据集合,还有很多的公开的开源的数据集合,比如 open catalyst 20( 到22)这样的数据集合。这些数据集合,它的来源背景实际上非常的不一样。比如老的 DP Library 中很多的数据集都是大家用来构建深度势能的数据集。而我们针对催化,材料等等领域,这些数据集实际上是另一种的构建思路——针对特定的领域,我去搜集尽可能多的数据。它们也有具有不同的 label 。他们使用不同的第一性原理的计算的软件,泛函、赝势,机组,有各种各样的区别,使得这些数据并不可能给我们一把就薅过来,一起训练也是不可能的。当然它们也有不同的复杂度。这是数据侧。在应用侧,我们面临的应用也是多种多样,上到天体物质模拟,下到分子制药等等。现在的模式,我们专对一个应用去特定的去产生一部分数据,训练一个模型,解决一个应用的问题。这样的方式相比于传统的用第一性原理计算,或者是用经验势去做,相比起来,无论是从精度和效率上来讲,其实已经有了非常大的提高。但是我们认为这个提高还没有到头,还仍然有空间。这里面其实是有一个终极的解决方案,有一个统一的模型。理论上来说,这些数据产生的背后的机理都是一样的,都是电子结构的计算。从根本原理上来讲,我们是能够训出一个统一的模型,去解决所有的问题。但是实际操作中,我们只能把它当作一个我们的终极目标。我们现阶段能够达到的能力,显然并不能够去生成这样的一个统一的模型,但是我们可以去要求我们的模型具备一定的条件,使得我们从数据到应用的 gap 缩的更加的小。
这个模型需要具有什么样的性质?首先,我的建模开销要与我的元素种类无关。因为我们希望从非常广泛的数据集里提取一些这个知识。大家一会可以看到一张图来说这件事情。第二,我们希望能够从数据集中提取构型和元素的知识。在数据提供来之后,这些数据并不是零散、无关的。首先他们最显然的关系就是元素周期表。这是大家从化学中总结出来的元素之间的关系。当然,我们从第一性原理的数据中,除了提取像元素周期表这种的数据以外,它们在内涵中也会有很多别的相关的知识,使得如果我们能够很好地去从中提炼这些知识,我们训练我们的模型所需要的数据量,就有可能并不是这些元素每个元素的数据量乘以元素个数这么多的数据,而是能够极大的得到降低。这就是为什么我们希望我们的模型能够很好的去从数据中提取一些知识,来降低模型训练对数据的需求。我们希望这样的模型能够达到一个预训练模型的效果。所谓预训练模型,意思就是我们从左侧的这些数据集中,能够达到我们前面说的这几点,这个模型就能够 encode 已有的数据中很大范围的知识。
当我们拿来一个特定的应用的时候,我们通过对预训练模型的特化,有希望能够使用比 from scratch 少很多的训练的开销,而获得一个精度上几乎可以媲美的模型。这就是我们预训练模型的构建思路。传统的DP模型实际上是并不满足这几点的,因此我们提出了基于注意力机制的DP模型。下面会给大家分享一些细节,我们做了什么样的事情,以及达到了什么样的效果。首先是在模型的设计方面,我们的出发点还是非常经典的DP模型里的环境矩阵,整个的构造方式是几乎一样的。不一样的地方就是我们引入了 Atomic Type Embedding。我们把原子元素的类型作为一个 one-hot (独热编码) 表示,再进行一个可学的映射,得到一个元素类型的空间的表示。这个表示是作为输入放在 Embedding Nets 里边的。值得注意是的,这里的Embedding Nets只有一个,但是传统的 DP 模型中, Embedding Nets 是依赖于中心原子以及邻居原子,或者是只依赖于邻居原子。在任一情形下,建模开销都是跟数据集里面的元素个数相关的。前者是 n square 的依赖,后者是 n 的依赖。如果我们选择一个 n square依赖的模型就会发现,随着元素种类的增长,整个模型的开销是增长趋势。显然,在这种趋势下,我们不可能用老的DP方式去建一个 50 多种元素的模型。但是如果在新的DPA-1模式下,我们做的这种 Type Embedding的模型的开销,是与元素个数增长是无关的,这使得训练几十个元素的大模型在这种设置下成为可能。
另外一个非常重要的点是,我们实际上是借鉴了已有的人工智能领域里面,已被证明非常有效的制作预训练模型的结构——Self-Attention 。首先对 Embedding Nets 的输出,是对每一个I原子,它是所有的邻居我们看可以把它看成一个序列,对这个序列做 Self-Attention。这个Q、K、V的构造都是非常标准的。不标准的地方就是 attention weights (注意力权重) 这个地方。attention weights 前半部分 SUBMAX 仍然是一个标准的构造,但是后边我们把它乘上了邻区里面角度的信息。这个部分加入,实际上是对模型表示能力的提升是非常关键的。attention weights 乘以 value 得到attention,我们可以对 embedding 的输出,做若干次的更新。实际中我们发现是两次就够。当然这里面需要做特定格式的 normalization 。后面是非常标准的 DP 构造,做出descriptor,然后再加上元素的 type 信息,经过 Fitnet 得到每个原子的能量贡献,后边都是跟标准的 DP 是一模一样。
在效果上面,除了刚才提到的模型开销和元素的个数无关以外,我们想强调的是,这个模型确实是学到了一些元素的本质的知识,使得它能够获得更好的精度。比如一个简单的例子,就是镁铝酮这样一个例子。我们的训练集是里面所有的单组元跟两组元的合金。我们的验证集不在训练集里面,验证集是三元的组分。我们就会发现,在这样的设置下,DPA-1 的精度是远远好于标准 DP 模型。看上去只学单组元跟两组元,标准 DP 实际上是不太知道三组元合金里面会发生什么,DPA-1实际上是已经通过单组元跟两组元的学习,知道了很多这些元素的构型和化学知识,在三组元的合金预测上已经达到了非常好的精度。这个 all 就是将所有的作为训练集。可以看到,相比于 baseline 来讲,其实它已经离得不太远。比如它已经学到了三组元里面的绝大多数的值。
标准的 DP 模型的能量的误差实际上是要大了一个量级,受力也是大了一倍左右。还有一些更极端的例子,比如 HEA,它是一个 6 元的合金,我们可以训,比如 interior 就是我们组分单纯型的内部, exterior 就是组分单纯型的外部。内部的是各组元比较均匀的占比,外部的就是某一个组元或者某几个组元,它的占比会非常高。我们可以看到, DPA-1只学内部就可以在外部上面有很好的精度,反之亦然。实际上我们可以看到它跟机械水平其实已经几乎是一样了,所以它只是学了一部分的元素的组成空间,基本上已经把所有 HEA 体系的知识都已经学会了。但是同样的事情,对于标准的 DP 模型来讲,它是很难做到的。
我们还考察了数据的效率,这里面所有的蓝线是标准的DP(模型),橘黄色的线就是我们的DPA-1的模型。横轴是 number of samples,纵轴是误差。大家可以看到,训练到相同的精度,我们的DPA-1模型相比于标准DP模型来讲,能够省下至多一个量级。这样的训练数据量说明它确实在数据效率的方面是非常高的。这里面还非常有意思的是测试绿线,我们拿OC 20 的这样的一个数据去做了 pre training,比如再以模型为基础进行特化的训练,得到绿线。虽然在现有的测试上面,它相比于橘红色线提升,并没有蓝线到红线黄线提升那么多,但是大家也注意到它是个 log scale,实际上数据量已经有本质性的降低了。但是我们需要指出的是,预训练再到特化的过程,我们现在用的方式仅是其中的一种,我们还在开发更好的方式,能够期望得到更好的预训练的效果。最后给大家分享一些更有意思的观察,就是模型通过我们新的这种构造方式,甚至具有了一些可解释的能力。我们的训练集是 OC20,它里面是一堆金属的substrate,上面放一些小分子的活性分子,在催化表面上会发生一些反应。这样的一个数据集,我们把它的元素嵌入的空间拿过来做一个PCA,选取最显著的三个维度,把所有里面 50 多种元素画出来之后,会发现它基本上给了我们元素周期表的这样的一个形状。但是它是一个螺旋性的元素周期表,沿着螺旋走是元素周期表的横向,就是同一层里面去填电子。垂直于螺旋的法向是元素周期表的纵向,我们电子层数再往上增加。这里面有些outliers,比如碳、氮、氧这些都是组成小分子的元素,所以它并不在螺旋上。其实也可以理解,因为这些数据在整个数据集里边的重要性,或者是构形是非常不一样的。这里面还有一个更有趣的观察,就是 OC20 里面,在碱金属这一组上面,它是没有锂的。如果它有锂,它应该是坐落在氢跟钠之间。我们就可以做一个非常有意思的考察,就是我们对氢跟钠的隐空间的表示,做一个差值,看看能不能recovery锂。事实表明它是能够一定程度上 recovery 锂。我们怎么考察?我们看差值的系数,在 1 的时候是钠,在 0 的时候是氢,纵轴是它在一个独立的锂的测试集上面去直接去进行测试的精度。这基本就意味着如果锂在表示之中,那里应该大概是在 0. 7 个钠加上 0. 3 个氢的这样的一个位置,确实是说明我们做了这样的一个 embedding 以后,不仅是形状上看着像元素周期表,确实是也能够在这样的一个隐空间的表示中,正确的获取未知元素的预测。这就标志着我们 DPA-1 的模型里面确实是具有一定的可解释性。
王涵,现为北京应用物理与计算数学研究所研究员,2002年进入北京大学数学科学学院学习,2006年获学士学位,2011年获博士学位,师从张平文院士。主要研究兴趣为分子模拟中的多尺度建模与计算方法,以及基于深度学习的分子建模。曾获中国数学会计算数学分会第五届青年创新奖。2020年与合作者共同推动完成的工作成果—“Pushing the limit of molecular dynamics with ab initio accuracy to 100 million atoms with machine learning”获得国际高性能计算应用领域最高奖—戈登贝尔奖(ACM Gordon Bell Prize)。